from torch import nn
import torch

"""
embed_dim–模型的总维度。(封装多头注意力机制中的query_size)
num_heads – 并行注意力头的数量。请注意，将被拆分 across （即每个 head 都将具有 dimension ）。embed_dimnum_headsembed_dim // num_heads
dropout – 的 的 dropout 概率。默认值：（无 dropout）。attn_output_weights0.0
bias – 如果指定，则向输入/输出投影图层添加偏差。Default: True.
add_bias_kv –如果指定，则在 dim=0 时向键和值序列添加偏差。Default:False.
add_zero_attn – 如果指定，则在 dim=1 处向键和值序列添加新的一批零。 Default:False
kdim – 键的特征总数。Default: None (uses kdim=embed_dim). (key_size)
vdim – 值的特征总数。Default: None (uses vdim=embed_dim).(value_size)
batch_first – 如果 ，则提供输入和输出张量 as （batch， seq， feature） 的 AS （Batch， seq， feature） 中。默认值：（seq， batch， feature）。TrueFalse
"""
mha = nn.MultiheadAttention(
    embed_dim=20,
    num_heads=4,
    dropout=0.5,
    bias=True,  # 最终通过w_o 输出的结果带有偏置
    add_bias_kv=True,  # kv相乘后的矩阵再添加纵向维度的一行数据作为偏置（这种添加对于最后的结果形态没有任何影响）
    add_zero_attn=False,
    kdim=20,
    vdim=20,
    batch_first=True
)

if __name__ == '__main__':
    """
    query (Tensor) –将查询和一组键值对映射到输出。 有关更多详细信息，请参阅“Attention Is All You Need”。
    key (Tensor) –将查询和一组键值对映射到输出。 有关更多详细信息，请参阅“Attention Is All You Need”。 
    value (Tensor) –将查询和一组键值对映射到输出。 有关更多详细信息，请参阅“Attention Is All You Need”。 
    key_padding_mask (Optional[Tensor]) –如果提供，则 key 中指定的 padding 元素将 被关注忽略。当给定一个二进制掩码且值为 True 时， Attention 层上对应的值将被忽略。
    need_weights (bool) –output attn_output_weights。
    attn_mask (Optional[Tensor]) –阻止对某些位置的注意的 2D 或 3D 蒙版。将向所有人广播 2D 蒙版 批次，而 3D 掩码允许为每个批次的条目指定不同的掩码。
    """
    # 自注意力机制
    queries = keys = values = torch.randn(10, 5, 20)
    valid_lens = torch.arange(0, 5, 1)

    # PAD 或者 <pad> 或者 [PAD] 掩码 (注意，是每个句子的pad掩码 总共有10个句子)
    key_padding_mask = torch.randint(0, 1, (10, 5)).float()  # 大于0的值认定为padding不参与于运算
    # 注意力机制掩码 （因果掩码）
    attn_mask = nn.Transformer.generate_square_subsequent_mask(5)
    results, weights = mha(
        query=queries,
        key=keys,
        value=values,
        key_padding_mask=key_padding_mask,
        need_weights=True,
        attn_mask=attn_mask
    )
    print(results.shape)
